{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0b8087bd",
   "metadata": {},
   "source": [
    "# Using MCP Toolbox with LlamaIndex\n",
    "\n",
    "Integrate your databases with LlamaIndex agents using MCP Toolbox."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a835390f",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "[MCP Toolbox for Databases](https://github.com/googleapis/genai-toolbox) is an open source MCP server for databases. It was designed with enterprise-grade and production-quality in mind. It enables you to develop tools easier, faster, and more securely by handling the complexities such as connection pooling, authentication, and more.\n",
    "\n",
    "Toolbox Tools can be seamlessly integrated with LlamaIndex applications. For more information on [getting started](https://googleapis.github.io/genai-toolbox/getting-started/local_quickstart/) or [configuring](https://googleapis.github.io/genai-toolbox/getting-started/configure/) MCP Toolbox, see the [documentation](https://googleapis.github.io/genai-toolbox/getting-started/introduction/).\n",
    "\n",
    "![architecture](https://raw.githubusercontent.com/googleapis/genai-toolbox/refs/heads/main/docs/en/getting-started/introduction/architecture.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4feb3728",
   "metadata": {},
   "source": [
    "## Configure and deploy\n",
    "\n",
    "Toolbox is an open source server that you deploy and manage yourself. For more\n",
    "instructions on deploying and configuring, see the official Toolbox\n",
    "documentation:\n",
    "\n",
    "* [Installing the Server](https://googleapis.github.io/genai-toolbox/getting-started/introduction/#installing-the-server)\n",
    "* [Configuring Toolbox](https://googleapis.github.io/genai-toolbox/getting-started/configure/)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38e6ade5",
   "metadata": {},
   "source": [
    "### Install client SDK\n",
    "\n",
    "Install the LlamaIndex compatible MCP Toolbox SDK package before getting started:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "305d03e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install toolbox-llamaindex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "966b587b",
   "metadata": {},
   "source": [
    "### Loading Toolbox Tools\n",
    "\n",
    "Once your Toolbox server is configured and up and running, you can load tools from your server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28c47c62",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ToolMetadata(description='Book a hotel by its ID. If the hotel is successfully booked, returns a NULL, raises an error if not.\\n\\nArgs:\\n    hotel_id (str): The ID of the hotel to book.', name='book-hotel', fn_schema=<class 'toolbox_core.utils.book-hotel'>, return_direct=False)\n",
      "ToolMetadata(description='Cancel a hotel by its ID.\\n\\nArgs:\\n    hotel_id (str): The ID of the hotel to cancel.', name='cancel-hotel', fn_schema=<class 'toolbox_core.utils.cancel-hotel'>, return_direct=False)\n",
      "ToolMetadata(description='Search for hotels based on location.\\n\\nArgs:\\n    location (str): The location of the hotel.', name='search-hotels-by-location', fn_schema=<class 'toolbox_core.utils.search-hotels-by-location'>, return_direct=False)\n",
      "ToolMetadata(description='Search for hotels based on name.\\n\\nArgs:\\n    name (str): The name of the hotel.', name='search-hotels-by-name', fn_schema=<class 'toolbox_core.utils.search-hotels-by-name'>, return_direct=False)\n",
      "ToolMetadata(description=\"Update a hotel's check-in and check-out dates by its ID. Returns a message indicating  whether the hotel was successfully updated or not.\\n\\nArgs:\\n    hotel_id (str): The ID of the hotel to update.\\n    checkin_date (str): The new check-in date of the hotel.\\n    checkout_date (str): The new check-out date of the hotel.\", name='update-hotel', fn_schema=<class 'toolbox_core.utils.update-hotel'>, return_direct=False)\n",
      "\n",
      "---- Find hotels in Basel with Basel in it's name. ----\n",
      "OK. I found three hotels in Basel with Basel in the name: Holiday Inn Basel, Hilton Basel, and Hyatt Regency Basel.\n",
      "\n",
      "\n",
      "---- Can you book the Hilton Basel for me? ----\n",
      "OK. I have booked the Hilton Basel for you.\n",
      "\n",
      "\n",
      "---- Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead. ----\n",
      "OK. I have booked the Hyatt Regency Basel for you.\n",
      "\n",
      "\n",
      "---- My check in dates would be from April 10, 2024 to April 19, 2024. ----\n",
      "OK. I have updated your check-in date to April 10, 2024 and your check-out date to April 19, 2024 for the Hyatt Regency Basel.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import os\n",
    "from llama_index.core.agent.workflow import AgentWorkflow\n",
    "from llama_index.core.workflow import Context\n",
    "from llama_index.llms.google_genai import GoogleGenAI\n",
    "from toolbox_llamaindex import ToolboxClient\n",
    "\n",
    "prompt = \"\"\"\n",
    "  You're a helpful hotel assistant. You handle hotel searching, booking and\n",
    "  cancellations. When the user searches for a hotel, mention it's name, id,\n",
    "  location and price tier. Always mention hotel ids while performing any\n",
    "  searches. This is very important for any operations. For any bookings or\n",
    "  cancellations, please provide the appropriate confirmation. Be sure to\n",
    "  update checkin or checkout dates if mentioned by the user.\n",
    "  Don't ask for confirmations from the user.\n",
    "\"\"\"\n",
    "\n",
    "queries = [\n",
    "    \"Find hotels in Basel with Basel in it's name.\",\n",
    "    \"Can you book the Hilton Basel for me?\",\n",
    "    \"Oh wait, this is too expensive. Please cancel it and book the Hyatt Regency instead.\",\n",
    "    \"My check in dates would be from April 10, 2024 to April 19, 2024.\",\n",
    "]\n",
    "\n",
    "\n",
    "async def run_application():\n",
    "    llm = GoogleGenAI(\n",
    "        api_key=os.getenv(\"GOOGLE_API_KEY\"),\n",
    "        model=\"gemini-2.0-flash-001\",\n",
    "    )\n",
    "\n",
    "    # llm = GoogleGenAI(\n",
    "    #     model=\"gemini-2.0-flash-001\",\n",
    "    #     vertexai_config={\"project\": \"project-id\", \"location\": \"us-central1\"},\n",
    "    # )\n",
    "\n",
    "    # Load the tools from the Toolbox server\n",
    "    async with ToolboxClient(\"http://127.0.0.1:5000\") as client:\n",
    "        tools = await client.aload_toolset()\n",
    "\n",
    "        agent = AgentWorkflow.from_tools_or_functions(\n",
    "            tools,\n",
    "            llm=llm,\n",
    "        )\n",
    "\n",
    "        for tool in tools:\n",
    "            print(tool.metadata)\n",
    "\n",
    "        ctx = Context(agent)\n",
    "\n",
    "        for query in queries:\n",
    "            response = await agent.run(user_msg=query, ctx=ctx)\n",
    "            print()\n",
    "            print(f\"---- {query} ----\")\n",
    "            print(str(response))\n",
    "\n",
    "\n",
    "await run_application()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffcfb58e",
   "metadata": {},
   "source": [
    "### Advanced Toolbox Features\n",
    "\n",
    "Toolbox has a variety of features to make developing Gen AI tools for databases.\n",
    "For more information, read more about the following features:\n",
    "\n",
    "* [Authenticated Parameters](https://googleapis.github.io/genai-toolbox/resources/tools/#authenticated-parameters): bind tool inputs to values from OIDC tokens automatically, making it easy to run sensitive queries without potentially leaking data\n",
    "* [Authorized Invocations](https://googleapis.github.io/genai-toolbox/resources/tools/#authorized-invocations): restrict access to use a tool based on the users Auth token\n",
    "* [OpenTelemetry](https://googleapis.github.io/genai-toolbox/how-to/export_telemetry/): get metrics and tracing from Toolbox with OpenTelemetry"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "senseAIenv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
